//
// Ephi - simulation of magnetic fields and particles
// Copyright (C) 2007 Indrek Mandre <indrek(at)mare.ee>
// For more information please see http://www.mare.ee/indrek/ephi/
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License along
// with this program; if not, write to the Free Software Foundation, Inc.,
// 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
//

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <assert.h>

#include "ephi.hpp"

#include <vector>

using namespace std;

#define RADIUS 0.15
#define SPACING 0.01
#define WR 0.035
#define CURRENT 50000
#define CHARGEPM 3e-7
#define DELTA_T 1e-11

#define STEPS 5000

struct particle
{
  Statics& statics;
  vect3d pos;
  vect3d v;
  prec_t charge_by_mass;
  prec_t delta_t;
  vect3d last_ef;
  vect3d last_bf;

  particle (Statics& statics, const vect3d& pos, const vect3d& v, prec_t charge, prec_t mass, prec_t delta_t) :
      statics(statics), pos(pos), v(v), charge_by_mass(charge/mass), delta_t(delta_t)
  {
    statics.getFields (pos, last_bf, last_ef);
  }

  void step ()
  {
    prec_t hdt = delta_t * 0.5;

    //statics.getFields (pos, last_bf, last_ef);

    vect3d k1 = charge_by_mass * (last_ef + v % last_bf);
    vect3d bf, ef;
    statics.getFields (pos + hdt * v + (0.125 * delta_t * delta_t) * k1, bf, ef);
    vect3d k2 = charge_by_mass * (ef + (v + hdt * k1) % bf);
    vect3d k3 = charge_by_mass * (ef + (v + hdt * k2) % bf);
    statics.getFields (pos + delta_t * v + (0.5 * delta_t * delta_t) * k3, last_bf, last_ef);
    vect3d k4 = charge_by_mass * (last_ef + (v + delta_t * k3) % last_bf);

    pos = pos + delta_t * v + (delta_t * delta_t / 6.0) * (k1 + k2 + k3);
    v = v + (delta_t / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4);
  }

  void astep ()
  {
    prec_t dt = delta_t;
    prec_t hdt = delta_t * 0.5;

    size_t count = 1;
    bool adapt = true;

    while (count)
      {
        //statics.getFields (pos, last_bf, last_ef);

        vect3d k1 = charge_by_mass * (last_ef + v % last_bf);
        vect3d bf, ef;
        statics.getFields (pos + hdt * v + (0.125 * dt * dt) * k1, bf, ef);
        vect3d k2 = charge_by_mass * (ef + (v + hdt * k1) % bf);

        if ( adapt )
          {
            prec_t k1len = k1.magnitude();
            prec_t kdiff = (k1 - k2).magnitude();
            if ( kdiff / k1len > 1e-3 * 1e-3 )
              {
                count *= 2;
                dt = delta_t / count;
                hdt = dt * 0.5;
                if ( count == 128 )
                  {
                    adapt = false;
                  }
                continue;
              }
            else
              {
                adapt = false;
              }
          }

        vect3d k3 = charge_by_mass * (ef + (v + hdt * k2) % bf);
        statics.getFields (pos + dt * v + (0.5 * dt * dt) * k3, last_bf, last_ef);
        vect3d k4 = charge_by_mass * (last_ef + (v + dt * k3) % last_bf);

        pos = pos + dt * v + (dt * dt / 6.0) * (k1 + k2 + k3);
        v = v + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4);

        count--;
      }
  }
};

struct entry
{
  vect3d pos;
  vect3d v;
  prec_t pdiff;
  prec_t vdiff;
  prec_t vdiffperc;
};

int main (int argc, char *argv[])
{
  Statics statics;
  make_polywell_cube (statics, RADIUS, WR, SPACING, CURRENT, CHARGEPM);

  std::vector<entry> POS;
  POS.reserve (STEPS);

  particle path(statics,
      vect3d(-2*RADIUS, 0.01, 0),
      vect3d(0, 0, 0),
      -ELEMENTARY_CHARGE,
      ELECTRON_MASS,
      DELTA_T);

  prec_t max_v = 0;
  prec_t cum_pdiff = 0;
  prec_t max_pdiff = 0;
  prec_t min_pdiff = 9e99;
  prec_t cum_vdiff = 0;
  prec_t max_vdiff = 0;
  prec_t cum_vdiffperc = 0;
  prec_t min_vdiffperc = 9e99;
  prec_t max_vdiffperc = 0;

  int n = 0;
  while (1)
    {
      particle p1(statics, path.pos, path.v, -ELEMENTARY_CHARGE, ELECTRON_MASS, DELTA_T);
      particle p2(statics, path.pos, path.v, -ELEMENTARY_CHARGE, ELECTRON_MASS, DELTA_T / 128);
      path.step();
      p1.astep();
      for ( int i = 0; i < 128; i++ )
          p2.step();
      prec_t pdiff = (p1.pos - p2.pos).length();
      prec_t vdiff = (p1.v - p2.v).length();
      prec_t vdiffperc = 100.0 * vdiff / p1.v.length();
#if 0
      printf ("%u: p.diff: %.3e, v.diff: %.3e (%.2e%%)\n",
          n, prec2double(pdiff), prec2double(vdiff), prec2double(vdiffperc));
#endif

      if ( n > 1500 ) {
      if ( min_pdiff > pdiff )
          min_pdiff = pdiff;
      if ( max_pdiff < pdiff )
          max_pdiff = pdiff;
      if ( max_vdiff < vdiff )
          max_vdiff = vdiff;
      if ( min_vdiffperc > vdiffperc )
          min_vdiffperc = vdiffperc;
      if ( max_vdiffperc < vdiffperc )
          max_vdiffperc = vdiffperc;
      }

      cum_pdiff += pdiff;
      cum_vdiff += vdiff;
      cum_vdiffperc += vdiffperc;
      if ( max_v < p1.v.length() )
          max_v = p1.v.length();

      entry e;
      e.pos = p1.pos;
      e.v = p1.v;
      e.pdiff = pdiff;
      e.vdiff = vdiff;
      e.vdiffperc = vdiffperc;
      POS.push_back (e);
      n++;
      if ( n >= STEPS )
          break;
    }

  printf ("max_v=%.2f\n", prec2double (max_v));
  printf ("max_pdiff=%g, avg_pdiff=%g\n", prec2double(max_pdiff), prec2double(cum_pdiff / n));
  printf ("max_vdiff=%g, avg_vdiff=%g\n", prec2double(max_vdiff), prec2double(cum_vdiff / n));
  printf ("max_vdiffperc=%g, avg_vdiffperc=%g\n", prec2double(max_vdiffperc), prec2double(cum_vdiffperc / n));

  // preparing the color gradient
  vector<Screen::color> colors;
  colors.reserve (256);
  for ( size_t i = 0; i <= 255; i++ )
    {
      Screen::color c;
      c.r = i;
      c.g = 255 - i;
      c.b = 0;
      colors.push_back (c);
    }

  Screen vscr(800, 800, RADIUS * 4.05, RADIUS * 4.05);
  vscr.draw_axes ();
  statics.draw (vscr);

  Screen pscr(800, 800, RADIUS * 4.05, RADIUS * 4.05);
  pscr.draw_axes ();
  statics.draw (pscr);

  for ( size_t i = 0; i < POS.size(); i++ )
    {
      entry& e = POS[i];
      size_t vindex = (size_t)floor(256.0 * (e.vdiffperc - min_vdiffperc) / (max_vdiffperc - min_vdiffperc));
      if ( vindex >= colors.size() ) vindex = colors.size() - 1;
      size_t pindex = (size_t)floor(256.0 * (e.pdiff - min_pdiff) / (max_pdiff - min_pdiff));
      if ( pindex >= colors.size() ) pindex = colors.size() - 1;
      vscr.set_pixel (e.pos, colors[vindex]);
      pscr.set_pixel (e.pos, colors[pindex]);
    }

  vscr.write ("rktest_v.bmp");
  pscr.write ("rktest_p.bmp");

  return 0;
}

